import tensorflow as tf
from PyQt5.QtCore import QThread, pyqtSignal

# 加载数据集线程
class LoadDataThread(QThread):
    # 设置触发信号传递的参数数据类型
    signal = pyqtSignal(str, list)
    def __init__(self, ui):
        super(LoadDataThread, self).__init__()
        self.train_dir = ""
        self.ui = ui

    def run(self):
        # 1. 加载数据集
        train_dataset, validate_dataset, class_names = self.m_data_load(self.train_dir, 224, 224, 16)
        self.ui.train_dataset = train_dataset
        self.ui.validate_dataset = validate_dataset
        # 2. 加载模型
        model = self.model_load(class_num=len(class_names))
        self.ui.model = model

        self.signal.emit(str(len(class_names)), class_names)

    # 模型加载
    def model_load(IMG_SHAPE=(224, 224, 3), class_num=214):
        base_model = tf.keras.applications.MobileNetV2(input_shape=(224, 224, 3), include_top=False, weights='imagenet')
        base_model.trainable = False

        model = tf.keras.models.Sequential([
            tf.keras.layers.experimental.preprocessing.Rescaling(1. / 127.5, offset=-1, input_shape=(224, 224, 3)),
            base_model,
            tf.keras.layers.GlobalAveragePooling2D(),
            tf.keras.layers.Dense(class_num, activation='softmax')
        ])

        # 输出模型信息
        model.summary()
        model.compile(optimizer='adam', loss='categorical_crossentropy',
                      metrics=['accuracy'])
        return model

    # 加载数据集
    def m_data_load(self, data_dir, img_height, img_width, batch_size):
        train_dataset = tf.keras.preprocessing.image_dataset_from_directory(
            data_dir,
            label_mode='categorical',
            validation_split=0.2,
            subset="training",
            seed=123,
            color_mode="rgb",
            image_size=(img_height, img_width),
            batch_size=batch_size)

        validate_dataset = tf.keras.preprocessing.image_dataset_from_directory(
            data_dir,
            label_mode='categorical',
            validation_split=0.2,
            subset="validation",
            seed=123,
            color_mode="rgb",
            image_size=(img_height, img_width),
            batch_size=batch_size)

        class_names = train_dataset.class_names

        return train_dataset, validate_dataset, class_names
